#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库管理器
负责数据库表的创建、管理和维护
支持自动创建爬虫相关的数据库表
"""

import os
import sys
import logging
import pymysql
from typing import Dict, List, Any, Optional
from datetime import datetime

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import DB_CONFIG, LOG_CONFIG

# 配置日志
logging.basicConfig(
    level=getattr(logging, LOG_CONFIG["level"]),
    format=LOG_CONFIG["format"]
)
logger = logging.getLogger(__name__)


class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self, config: Optional[Dict] = None):
        self.config = config or DB_CONFIG
        self.table_schemas = self._get_table_schemas()
        logger.info("数据库管理器初始化完成")
    
    def get_connection(self):
        """获取数据库连接"""
        try:
            connection = pymysql.connect(
                host=self.config['host'],
                port=self.config['port'],
                user=self.config.get('username', self.config.get('user', 'root')),
                password=self.config['password'],
                database=self.config['database'],
                charset='utf8mb4',
                cursorclass=pymysql.cursors.DictCursor,
                autocommit=False
            )
            return connection
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            raise
    
    def _get_table_schemas(self) -> Dict[str, str]:
        """获取数据库表结构定义"""
        return {
            'news_data': """
                CREATE TABLE IF NOT EXISTS `news_data` (
                    `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '主键ID',
                    `url` varchar(500) DEFAULT NULL COMMENT '文章URL',
                    `title` varchar(500) DEFAULT NULL COMMENT '文章标题',
                    `content` longtext COMMENT '文章正文内容',
                    `author` varchar(100) DEFAULT NULL COMMENT '作者',
                    `publishTime` bigint(20) DEFAULT NULL COMMENT '发布时间戳',
                    `createTimeEs` varchar(50) DEFAULT NULL COMMENT '创建时间',
                    `publishChannel` varchar(10) DEFAULT NULL COMMENT '发布渠道',
                    `webName` varchar(100) DEFAULT NULL COMMENT '网站名称',
                    `columnName` varchar(100) DEFAULT NULL COMMENT '栏目名称',
                    `mediaName` varchar(100) DEFAULT NULL COMMENT '媒体名称',
                    `mediaOrganization` varchar(100) DEFAULT NULL COMMENT '媒体机构',
                    `simhash` bigint(20) DEFAULT NULL COMMENT 'Simhash值',
                    `attitude` varchar(20) DEFAULT NULL COMMENT '情感属性',
                    `newsLevel` varchar(20) DEFAULT NULL COMMENT '新闻重要级别',
                    `keyWord` varchar(200) DEFAULT NULL COMMENT '关键字',
                    `mediaType` varchar(10) DEFAULT NULL COMMENT '媒体类型',
                    `summary` text DEFAULT NULL COMMENT '文章摘要',
                    `category` varchar(100) DEFAULT NULL COMMENT '分类',
                    `tags` varchar(200) DEFAULT NULL COMMENT '标签',
                    `source_type` varchar(50) DEFAULT 'universal_spider' COMMENT '数据来源类型',
                    `created_at` timestamp DEFAULT CURRENT_TIMESTAMP COMMENT '记录创建时间',
                    `updated_at` timestamp DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '记录更新时间',
                    PRIMARY KEY (`id`),
                    UNIQUE KEY `idx_url` (`url`),
                    KEY `idx_title` (`title`),
                    KEY `idx_publish_time` (`publishTime`),
                    KEY `idx_media_name` (`mediaName`),
                    KEY `idx_category` (`category`),
                    KEY `idx_source_type` (`source_type`),
                    KEY `idx_created_at` (`created_at`)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='新闻数据表'
            """,
            
            'spider_schedule': """
                CREATE TABLE IF NOT EXISTS `spider_schedule` (
                    `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '主键ID',
                    `task_name` varchar(100) NOT NULL COMMENT '任务名称',
                    `site_name` varchar(100) NOT NULL COMMENT '网站名称',
                    `site_url` varchar(500) NOT NULL COMMENT '网站URL',
                    `column_name` varchar(100) DEFAULT NULL COMMENT '栏目名称',
                    `content_type` varchar(50) DEFAULT 'news' COMMENT '内容类型',
                    `tool_type` varchar(20) DEFAULT 'requests' COMMENT '爬取工具类型',
                    `use_ai` tinyint(1) DEFAULT 0 COMMENT '是否使用AI解析',
                    `max_links` int(11) DEFAULT 50 COMMENT '最大链接数',
                    `delay_min` int(11) DEFAULT 1 COMMENT '最小延迟秒数',
                    `delay_max` int(11) DEFAULT 3 COMMENT '最大延迟秒数',
                    `status` enum('0','1','2') DEFAULT '1' COMMENT '状态：0=禁用，1=启用，2=暂停',
                    `priority` int(11) DEFAULT 5 COMMENT '优先级：1-10，数字越小优先级越高',
                    `cron_expression` varchar(50) DEFAULT NULL COMMENT 'Cron表达式',
                    `last_run_time` timestamp NULL DEFAULT NULL COMMENT '上次运行时间',
                    `next_run_time` timestamp NULL DEFAULT NULL COMMENT '下次运行时间',
                    `run_count` int(11) DEFAULT 0 COMMENT '运行次数',
                    `success_count` int(11) DEFAULT 0 COMMENT '成功次数',
                    `error_count` int(11) DEFAULT 0 COMMENT '错误次数',
                    `last_error` text DEFAULT NULL COMMENT '最后一次错误信息',
                    `config_json` text DEFAULT NULL COMMENT '配置信息JSON',
                    `created_by` varchar(50) DEFAULT 'system' COMMENT '创建者',
                    `created_at` timestamp DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                    `updated_at` timestamp DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                    PRIMARY KEY (`id`),
                    UNIQUE KEY `idx_task_name` (`task_name`),
                    KEY `idx_site_name` (`site_name`),
                    KEY `idx_status` (`status`),
                    KEY `idx_priority` (`priority`),
                    KEY `idx_next_run_time` (`next_run_time`)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='爬虫调度任务表'
            """,
            
            'spider_logs': """
                CREATE TABLE IF NOT EXISTS `spider_logs` (
                    `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '主键ID',
                    `task_id` int(11) DEFAULT NULL COMMENT '任务ID',
                    `task_name` varchar(100) DEFAULT NULL COMMENT '任务名称',
                    `site_url` varchar(500) DEFAULT NULL COMMENT '网站URL',
                    `log_level` varchar(20) DEFAULT 'INFO' COMMENT '日志级别',
                    `log_type` varchar(50) DEFAULT 'spider' COMMENT '日志类型',
                    `message` text DEFAULT NULL COMMENT '日志消息',
                    `details` json DEFAULT NULL COMMENT '详细信息JSON',
                    `start_time` timestamp NULL DEFAULT NULL COMMENT '开始时间',
                    `end_time` timestamp NULL DEFAULT NULL COMMENT '结束时间',
                    `duration` int(11) DEFAULT NULL COMMENT '执行时长(秒)',
                    `status` varchar(20) DEFAULT 'running' COMMENT '状态',
                    `links_found` int(11) DEFAULT 0 COMMENT '发现链接数',
                    `articles_extracted` int(11) DEFAULT 0 COMMENT '提取文章数',
                    `articles_saved` int(11) DEFAULT 0 COMMENT '保存文章数',
                    `error_count` int(11) DEFAULT 0 COMMENT '错误次数',
                    `created_at` timestamp DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                    PRIMARY KEY (`id`),
                    KEY `idx_task_id` (`task_id`),
                    KEY `idx_task_name` (`task_name`),
                    KEY `idx_log_level` (`log_level`),
                    KEY `idx_log_type` (`log_type`),
                    KEY `idx_status` (`status`),
                    KEY `idx_created_at` (`created_at`)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='爬虫运行日志表'
            """,
            
            'spider_statistics': """
                CREATE TABLE IF NOT EXISTS `spider_statistics` (
                    `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '主键ID',
                    `date` date NOT NULL COMMENT '统计日期',
                    `task_name` varchar(100) DEFAULT NULL COMMENT '任务名称',
                    `site_name` varchar(100) DEFAULT NULL COMMENT '网站名称',
                    `total_runs` int(11) DEFAULT 0 COMMENT '总运行次数',
                    `success_runs` int(11) DEFAULT 0 COMMENT '成功运行次数',
                    `failed_runs` int(11) DEFAULT 0 COMMENT '失败运行次数',
                    `total_links` int(11) DEFAULT 0 COMMENT '总链接数',
                    `total_articles` int(11) DEFAULT 0 COMMENT '总文章数',
                    `total_duration` int(11) DEFAULT 0 COMMENT '总耗时(秒)',
                    `avg_duration` decimal(10,2) DEFAULT 0.00 COMMENT '平均耗时(秒)',
                    `success_rate` decimal(5,2) DEFAULT 0.00 COMMENT '成功率(%)',
                    `created_at` timestamp DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                    `updated_at` timestamp DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                    PRIMARY KEY (`id`),
                    UNIQUE KEY `idx_date_task` (`date`, `task_name`),
                    KEY `idx_date` (`date`),
                    KEY `idx_site_name` (`site_name`)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='爬虫统计表'
            """,
            
            'ai_parse_logs': """
                CREATE TABLE IF NOT EXISTS `ai_parse_logs` (
                    `id` int(11) NOT NULL AUTO_INCREMENT COMMENT '主键ID',
                    `article_url` varchar(500) DEFAULT NULL COMMENT '文章URL',
                    `provider_name` varchar(50) DEFAULT NULL COMMENT 'AI服务提供商',
                    `content_type` varchar(50) DEFAULT 'news' COMMENT '内容类型',
                    `input_length` int(11) DEFAULT 0 COMMENT '输入文本长度',
                    `output_length` int(11) DEFAULT 0 COMMENT '输出文本长度',
                    `fields_requested` json DEFAULT NULL COMMENT '请求提取的字段',
                    `fields_extracted` json DEFAULT NULL COMMENT '成功提取的字段',
                    `parse_status` varchar(20) DEFAULT 'success' COMMENT '解析状态',
                    `error_message` text DEFAULT NULL COMMENT '错误信息',
                    `response_time` int(11) DEFAULT 0 COMMENT '响应时间(毫秒)',
                    `token_usage` json DEFAULT NULL COMMENT 'Token使用情况',
                    `cost` decimal(10,4) DEFAULT 0.0000 COMMENT '成本',
                    `created_at` timestamp DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                    PRIMARY KEY (`id`),
                    KEY `idx_article_url` (`article_url`),
                    KEY `idx_provider_name` (`provider_name`),
                    KEY `idx_content_type` (`content_type`),
                    KEY `idx_parse_status` (`parse_status`),
                    KEY `idx_created_at` (`created_at`)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='AI解析日志表'
            """
        }
    
    def check_database_exists(self) -> bool:
        """检查数据库是否存在"""
        try:
            # 临时连接，不指定数据库
            temp_config = self.config.copy()
            database_name = temp_config.pop('database')
            
            conn = pymysql.connect(**temp_config)
            with conn.cursor() as cursor:
                cursor.execute("SHOW DATABASES LIKE %s", (database_name,))
                result = cursor.fetchone()
                conn.close()
                return result is not None
        except Exception as e:
            logger.error(f"检查数据库存在性失败: {e}")
            return False
    
    def create_database(self) -> bool:
        """创建数据库"""
        try:
            # 临时连接，不指定数据库
            temp_config = self.config.copy()
            database_name = temp_config.pop('database')
            
            conn = pymysql.connect(**temp_config)
            with conn.cursor() as cursor:
                cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{database_name}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
                logger.info(f"数据库 {database_name} 创建成功")
                conn.close()
                return True
        except Exception as e:
            logger.error(f"创建数据库失败: {e}")
            return False
    
    def check_table_exists(self, table_name: str) -> bool:
        """检查表是否存在"""
        try:
            conn = self.get_connection()
            with conn.cursor() as cursor:
                cursor.execute("SHOW TABLES LIKE %s", (table_name,))
                result = cursor.fetchone()
                conn.close()
                return result is not None
        except Exception as e:
            logger.error(f"检查表 {table_name} 存在性失败: {e}")
            return False
    
    def create_table(self, table_name: str) -> bool:
        """创建单个表"""
        if table_name not in self.table_schemas:
            logger.error(f"未找到表 {table_name} 的结构定义")
            return False
        
        try:
            conn = self.get_connection()
            with conn.cursor() as cursor:
                cursor.execute(self.table_schemas[table_name])
                conn.commit()
                logger.info(f"表 {table_name} 创建成功")
                conn.close()
                return True
        except Exception as e:
            logger.error(f"创建表 {table_name} 失败: {e}")
            return False
    
    def create_all_tables(self) -> Dict[str, bool]:
        """创建所有表"""
        results = {}
        
        for table_name in self.table_schemas.keys():
            if self.check_table_exists(table_name):
                logger.info(f"表 {table_name} 已存在，跳过创建")
                results[table_name] = True
            else:
                results[table_name] = self.create_table(table_name)
        
        return results
    
    def initialize_database(self) -> bool:
        """初始化数据库（创建数据库和所有表）"""
        logger.info("开始初始化数据库...")
        
        # 检查并创建数据库
        if not self.check_database_exists():
            logger.info("数据库不存在，开始创建...")
            if not self.create_database():
                logger.error("数据库创建失败")
                return False
        else:
            logger.info("数据库已存在")
        
        # 创建所有表
        results = self.create_all_tables()
        
        # 检查结果
        success_count = sum(1 for success in results.values() if success)
        total_count = len(results)
        
        if success_count == total_count:
            logger.info(f"数据库初始化完成，成功创建/验证 {success_count} 个表")
            return True
        else:
            logger.error(f"数据库初始化部分失败，{success_count}/{total_count} 个表创建成功")
            return False
    
    def get_table_info(self, table_name: str) -> dict:
        """获取表信息"""
        try:
            conn = self.get_connection()
            
            with conn.cursor() as cursor:
                # 检查表是否存在
                cursor.execute("""
                    SELECT COUNT(*) 
                    FROM information_schema.tables 
                    WHERE table_schema = %s AND table_name = %s
                """, (self.config['database'], table_name))
                
                table_exists = cursor.fetchone()[0] > 0
                
                if not table_exists:
                    return {
                        'exists': False,
                        'record_count': 0,
                        'columns': []
                    }
                
                # 获取记录数
                cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
                record_count = cursor.fetchone()[0]
                
                # 获取列信息
                cursor.execute(f"DESCRIBE {table_name}")
                columns = cursor.fetchall()
                
            conn.close()
            
            return {
                'exists': True,
                'record_count': record_count,
                'columns': columns
            }
            
        except Exception as e:
            logging.error(f"获取表 {table_name} 信息失败: {e}")
            return {
                'exists': False,
                'record_count': 0,
                'columns': [],
                'error': str(e)
            }
    
    def get_database_status(self) -> Dict[str, Any]:
        """获取数据库状态"""
        status = {
            'database_exists': self.check_database_exists(),
            'tables': {},
            'total_tables': len(self.table_schemas),
            'existing_tables': 0,
            'total_records': 0
        }
        
        if status['database_exists']:
            for table_name in self.table_schemas.keys():
                table_info = self.get_table_info(table_name)
                status['tables'][table_name] = table_info
                
                if table_info['exists']:
                    status['existing_tables'] += 1
                    status['total_records'] += table_info.get('record_count', 0)
        
        return status
    
    def insert_sample_schedule_task(self) -> bool:
        """插入示例调度任务"""
        try:
            conn = self.get_connection()
            with conn.cursor() as cursor:
                # 检查是否已有数据
                cursor.execute("SELECT COUNT(*) FROM spider_schedule")
                count = cursor.fetchone()[0]
                
                if count == 0:
                    # 插入示例任务
                    sample_tasks = [
                        {
                            'task_name': 'CCTV新闻爬取',
                            'site_name': 'CCTV新闻',
                            'site_url': 'https://news.cctv.com/',
                            'column_name': '新闻频道',
                            'content_type': 'news',
                            'tool_type': 'requests',
                            'use_ai': 0,
                            'max_links': 20,
                            'status': '1',
                            'priority': 5
                        },
                        {
                            'task_name': '人民网新闻爬取',
                            'site_name': '人民网',
                            'site_url': 'http://www.people.com.cn/',
                            'column_name': '要闻',
                            'content_type': 'news',
                            'tool_type': 'requests',
                            'use_ai': 0,
                            'max_links': 30,
                            'status': '1',
                            'priority': 3
                        }
                    ]
                    
                    for task in sample_tasks:
                        sql = """
                        INSERT INTO spider_schedule (
                            task_name, site_name, site_url, column_name, content_type,
                            tool_type, use_ai, max_links, status, priority
                        ) VALUES (
                            %(task_name)s, %(site_name)s, %(site_url)s, %(column_name)s, %(content_type)s,
                            %(tool_type)s, %(use_ai)s, %(max_links)s, %(status)s, %(priority)s
                        )
                        """
                        cursor.execute(sql, task)
                    
                    conn.commit()
                    logger.info(f"插入 {len(sample_tasks)} 个示例调度任务")
                else:
                    logger.info("调度任务表已有数据，跳过示例数据插入")
                
                conn.close()
                return True
                
        except Exception as e:
            logger.error(f"插入示例调度任务失败: {e}")
            return False
    
    def optimize_tables(self) -> bool:
        """优化表（分析表和优化表）"""
        try:
            conn = self.get_connection()
            with conn.cursor() as cursor:
                for table_name in self.table_schemas.keys():
                    if self.check_table_exists(table_name):
                        # 分析表
                        cursor.execute(f"ANALYZE TABLE {table_name}")
                        # 优化表
                        cursor.execute(f"OPTIMIZE TABLE {table_name}")
                        logger.info(f"表 {table_name} 优化完成")
                
                conn.close()
                logger.info("所有表优化完成")
                return True
                
        except Exception as e:
            logger.error(f"优化表失败: {e}")
            return False


# 使用示例
def main():
    """主函数"""
    try:
        # 创建数据库管理器
        db_manager = DatabaseManager()
        
        # 初始化数据库
        if db_manager.initialize_database():
            print("✅ 数据库初始化成功")
            
            # 插入示例数据
            if db_manager.insert_sample_schedule_task():
                print("✅ 示例数据插入成功")
            
            # 获取数据库状态
            status = db_manager.get_database_status()
            print(f"\n📊 数据库状态:")
            print(f"  数据库存在: {status['database_exists']}")
            print(f"  总表数: {status['total_tables']}")
            print(f"  已存在表数: {status['existing_tables']}")
            print(f"  总记录数: {status['total_records']}")
            
            print("\n📋 表详情:")
            for table_name, table_info in status['tables'].items():
                if table_info['exists']:
                    print(f"  {table_name}: ✅ ({table_info['record_count']} 条记录)")
                else:
                    print(f"  {table_name}: ❌")
        else:
            print("❌ 数据库初始化失败")
            
    except Exception as e:
        logger.error(f"数据库管理器测试失败: {e}")


if __name__ == "__main__":
    main() 